import os
from glob import glob
import torch
import matplotlib.pyplot as plt
import numpy as np
import cv2
import math


def save_sample_imgs(sample_gt, sample_pre, save_dir, epoch, step):
    totalsaveitem = len(sample_gt.keys()) + len(sample_pre.keys())
    subitem = int(math.ceil(1.0*totalsaveitem/2.0))

    print("save sample to %s"%save_dir)
    pltindex = 1
    fig = plt.figure(figsize=(19, 10))

    image = sample_gt['image'][0].data.cpu().numpy()
    image = np.transpose(image, (1, 2, 0))
    image = 255 * (image * 0.5 + 0.5)
    image = image.astype(np.uint8)
    f = fig.add_subplot(2, subitem, pltindex)
    pltindex += 1
    f.imshow(image)
    f.set_title("image")

    for target_name in ['mask_gt', 'blurmask_gt', 'alpha_gt']:
        if target_name in sample_gt:
            mask_gt = sample_gt[target_name][0][0].data.cpu().numpy()
            mask_gt = 255 * mask_gt
            mask_gt = mask_gt.astype(np.uint8)
            f = fig.add_subplot(2, subitem, pltindex)
            pltindex += 1
            f.imshow(mask_gt)
            f.set_title(target_name)

    for target_name in ['mask_pred', 'blurmask_pred', 'alpha_pred', 'dyemask']:
        if target_name in sample_pre:
            targer_result = sample_pre[target_name].data.cpu().numpy()

            target_dim = len(targer_result.shape)

            if target_dim == 4:
                targer_result = targer_result[0].squeeze()

            if len(targer_result.shape) == 3:
                targer_result = targer_result.argmax(0)

            if target_name != 'dyemask':
                targer_result = 255 * targer_result
                targer_result = targer_result.astype(np.uint8)
                f = fig.add_subplot(2, subitem, pltindex)
                pltindex += 1
                f.imshow(targer_result)
                f.set_title(target_name)
            else:
                dyeimg = hairrecolor(image, targer_result, (0x40, 0x16, 0x66))
                f = fig.add_subplot(2, subitem, pltindex)
                pltindex += 1
                f.imshow(dyeimg)
                f.set_title("dye")

    p = os.path.join(save_dir, "epoch-%s_step-%s.png" % (epoch, step))
    plt.savefig(p)

def getimages(data_folder):
    for root, _, filenames in os.walk(data_folder):
        for filename in filenames:
            imgpath = os.path.join(root, filename)
            bgr_img = cv2.imread(imgpath)
            rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
            yield bgr_img, rgb_img



def showimg(image):
    image = image[0]
    image = image.cpu().detach().numpy()
    image = np.transpose(image,(1,2,0))
    image = 255*(image*0.5+0.5) 
    image = image.astype(np.uint8)

    cv2.imshow("image", image)
    cv2.waitKey(0)



def showmask(mask):
    mask = mask[0]
    mask = mask.cpu().detach().numpy()
    mask = mask[0]
    mask = 255*(mask)
    mask = mask.astype(np.uint8)

    cv2.imshow("mask", mask)
    cv2.waitKey(0)


def set_lr(args, epoch, optimizer):
    lrDecay = args.lrDecay
    decayType = args.lrdecayType
    if decayType == 'keep':
        lr = args.lr
    elif decayType == 'step':
        epoch_iter = (epoch + 1) // lrDecay
        lr = args.lr / 2 ** epoch_iter
    elif decayType == 'exp':
        k = math.log(2) / lrDecay
        lr = args.lr * math.exp(-k * epoch)
    elif decayType == 'poly':
        lr = args.lr * math.pow((1 - epoch / args.nEpochs), 0.9)

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    return lr



class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self, name="unknown"):
        self.is_use = False
        self.name = name
        self.reset()

    def reset(self):
        self.is_use = False
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.is_use = True
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def getFilepath(rootdir, listfile):
    image_list = []
    blurmask_list = []
    mask_list = []
    trimap_list = []
    alpha_list = []

    with open(listfile, "r") as f:
        for line in f.readlines():
            line = line.strip()
            infor = line.split(",")
            image_list.append(os.path.join(rootdir, infor[0]))
            blurmask_list.append(os.path.join(rootdir, infor[1]))
            mask_list.append(os.path.join(rootdir, infor[2]))
            trimap_list.append(os.path.join(rootdir, infor[3]))
            alpha_list.append(os.path.join(rootdir, infor[4]))

    assert(len(image_list) == len(mask_list) and len(trimap_list) == len(alpha_list) )
    return image_list, blurmask_list, mask_list, trimap_list, alpha_list


def np2Tensor(array):
    if len(array.shape) == 3:
        tensor = torch.FloatTensor(array.transpose((2, 0, 1)).astype(float))
    else:
        tensor = torch.FloatTensor(np.expand_dims(array, 0).astype(float))
    return tensor




def change_v(v, mask, target):
    # 染发
    epsilon = 1e-7
    x = v / 255                             # 数学化
    target = target / 255
    target = -np.log(epsilon + 1 - target)

    masksum = np.sum(mask)
    if masksum != 0.0:
        x_mean = np.sum(-np.log(epsilon + 1 - x)  * mask) / masksum
        alpha = target / x_mean
        x = 1 - (1 - x) ** alpha
        v[:] = x * 255                          # 二进制化
    else:
        v[:] = x * 255  # 二进制化


def hairrecolor(im, mask, color=(0x40, 0x16, 0x66)):
    h, w = im.shape[0:2]
    mask = cv2.resize(mask, (w,h))
    mask = mask[:, :, np.newaxis]
    # 工程化
    color = np.array(color, dtype='uint8', ndmin=3)
    im_hsv = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)
    color_hsv = cv2.cvtColor(color, cv2.COLOR_BGR2HSV)
    # 染发
    im_hsv[..., 0] = color_hsv[..., 0]      # 修改颜色
    change_v(im_hsv[..., 2:], mask, color_hsv[..., 2:])
    im_hsv[..., 1] = color_hsv[..., 1]      # 修改饱和度
    x = cv2.cvtColor(im_hsv, cv2.COLOR_HSV2BGR)
    im = im * (1 - mask) + x * mask
    im = im.astype(np.uint8)
    return im


def show_sample_withplt(sample):
    image = sample['image'].numpy()
    image = np.transpose(image, (1, 2, 0))
    image = 255 * (image * 0.5 + 0.5)
    image = image.astype(np.uint8)

    plt.get_current_fig_manager().full_screen_toggle()
    plt.subplot(231)
    plt.imshow(image)

    if 'mask' in sample:
        mask = sample['mask'].numpy()[0]
        mask = 255 * mask
        mask = mask.astype(np.uint8)
        plt.subplot(232)
        plt.imshow(mask)

    if 'blurmask' in sample:
        blurmask = sample['blurmask'].numpy()[0]
        blurmask = 255 * blurmask
        blurmask = blurmask.astype(np.uint8)
        plt.subplot(233)
        plt.imshow(blurmask)

    if 'trimap' in sample:
        trimap = sample['trimap'].numpy()[0]
        trimap = 0 * (trimap == 0) + 128 * (trimap == 1) + 255 * (trimap == 2)
        trimap = trimap.astype(np.uint8)
        plt.subplot(234)
        plt.imshow(trimap)

    if 'alpha' in sample:
        alpha = sample['alpha'].numpy()[0]
        alpha = 255 * alpha
        alpha = alpha.astype(np.uint8)
        plt.subplot(235)
        plt.imshow(alpha)

    if 'blurmask' in sample:
        alpha = sample['blurmask'].numpy()[0]
        dyeimg = hairrecolor(image, alpha, (0x40, 0x16, 0x66))
        plt.subplot(236)
        plt.imshow(dyeimg)


    plt.show()


def show_sample_withocv(sample):
    image = sample['image'].numpy()
    image = np.transpose(image, (1, 2, 0))
    image = 255 * (image * 0.5 + 0.5)
    image = image.astype(np.uint8)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    if 'mask_gt' in sample:
        mask = sample['mask_gt'].numpy()[0]
        mask = 255 * mask
        mask = mask.astype(np.uint8)
        cv2.imshow("mask_gt", mask)

    if 'blurmask_gt' in sample:
        blurmask = sample['blurmask_gt'].numpy()[0]
        blurmask = 255 * blurmask
        blurmask = blurmask.astype(np.uint8)
        cv2.imshow("blurmask_gt", blurmask)

    if 'trimap_gt' in sample:
        trimap = sample['trimap_gt'].numpy()[0]
        trimap = 0 * (trimap == 0) + 128 * (trimap == 1) + 255 * (trimap == 2)
        trimap = trimap.astype(np.uint8)
        cv2.imshow("trimap_gt", trimap)

    if 'alpha_gt' in sample:
        alpha = sample['alpha_gt'].numpy()[0]
        alpha = 255 * alpha
        alpha = alpha.astype(np.uint8)
        cv2.imshow("alpha_gt", alpha)

    if 'blurmask_gt' in sample:
        alpha = sample['blurmask_gt'].numpy()[0]
        dyeimg = hairrecolor(image, alpha, (0x40, 0x16, 0x66))
        cv2.imshow("dye", dyeimg)

    cv2.waitKey(0)

#
def batch_end_log(**kwargs):
    if kwargs['step'] % kwargs['log_intervel'] == 0:
        format_str = "epoch:[{epoch}/{max_epoch}] image:[{step}/{max_step}] ".format(**kwargs)
        for key, value in kwargs.items():
            if key not in ['epoch', 'max_epoch', 'step', 'max_step', 'log_intervel']:
                format_str += "%s:%f "%(key, value)
        print(format_str)


def epoch_end_log(**kwargs):

    format_str = "%s epoch done" % kwargs["phase"]

    for key, value in kwargs.items():
        if key not in ['phase']:
            format_str += "%s:%f "%(key, value)
    print(format_str)

def make_ckpt_name(**kwargs):

    format_str = ""
    for key, value in kwargs.items():
        format_str += "{}:{}_".format(key, value)
    return format_str[:-1]+"."+kwargs['suffix']


color_print_map={"red":31, "green":32, "blue":34, "white":37, "yellow":33}
def color_print(printstr, color="red"):
    print("\033[1;%d;40m%s\033[0m" % (color_print_map[color], printstr))
